import argparse

import cv2
import mmcv
import onnxruntime


def visualize_bbox(img, bbox, class_name, thickness=2):
    """Visualizes a single bounding box on the image"""
    BOX_COLOR = (255, 0, 0)  # Red
    TEXT_COLOR = (255, 255, 255)  # White
    if len(bbox) >= 5:
        x_min, y_min, x_max, y_max = [int(e) for e in bbox[:-1]]
        score = bbox[-1]
        str_info = F"{class_name} {score:.2f}"
    else:
        assert len(bbox) == 4
        x_min, y_min, x_max, y_max = [int(e) for e in bbox]
        str_info = F"{class_name}"
    cv2.rectangle(img, (x_min, y_min), (x_max, y_max),
                  color=BOX_COLOR,
                  thickness=thickness)

    ((text_width, text_height), _) = cv2.getTextSize(str_info,
                                                     cv2.FONT_HERSHEY_SIMPLEX,
                                                     0.35, 1)
    cv2.rectangle(img, (x_min, y_min - int(1.3 * text_height)),
                  (x_min + text_width, y_min), BOX_COLOR, -1)
    cv2.putText(
        img,
        text=str_info,
        org=(x_min, y_min - int(0.3 * text_height)),
        fontFace=cv2.FONT_HERSHEY_SIMPLEX,
        fontScale=0.35,
        color=TEXT_COLOR,
        lineType=cv2.LINE_AA,
    )
    return img


def argument_parser():
    parser = argparse.ArgumentParser(description="export model to onnx")
    parser.add_argument("--onnx-file",
                        default="out.onnx",
                        help="path to onnx file")
    parser.add_argument("--image",
                        default="images/person.jpg",
                        help="path to image file")
    args = parser.parse_args()
    return args


def main():
    args = argument_parser()
    onnx_file_path = args.onnx_file
    image_path = args.image

    image = mmcv.imread(image_path)

    model = onnxruntime.InferenceSession(onnx_file_path)
    input_name = model.get_inputs()[0].name
    output_names = [output.name for output in model.get_outputs()]

    cls_scores, bbox_preds = model.run(output_names,
                                       input_feed={input_name: image})
    for score, box in zip(cls_scores, bbox_preds):
        visualize_bbox(image, box.tolist(), "person")
    cv2.imshow("show", image)
    cv2.waitKey()
    cv2.destroyWindow("show")


if __name__ == "__main__":
    main()
